#
# messkit.simple_stream
#
# Copyright 2006-2008 Helsinki Institute for Information Technology (HIIT)
# and the authors.
#
# Authors: Ken Rimey <rimey@hiit.fi>
#

# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

"""
Stream implementation based on asyncore

This strictly single-threaded implementation does not use the
callback-callback mechanism.
"""

import os
import sys
import socket
import logging
import asyncore
from Queue import Queue, Empty

from messkit.stream import Stream as _Stream
from messkit.stream import StreamServer as _StreamServer
from messkit.stream import StreamManager as _StreamManager

__all__ = ['Stream', 'StreamServer', 'StreamManager']

class Channel(asyncore.dispatcher):
    def __init__(self, map, readable=None, writable=None):
        asyncore.dispatcher.__init__(self, map=map)

        # We have two queues of handlers.  Each handler is a
        # generator, which will yield True as long as it wants to
        # remain in effect.  This convention, where the handler can
        # either yield False or raise StopIteration to signal that it
        # is finished, makes it easier to write handlers that always
        # finish their work in the first call to next().
        self.read_handlers = []
        self.write_handlers = []

        # Extra predicates, useful for rate limiting:
        self._readable = readable or (lambda: True)
        self._writable = writable or (lambda: True)

    def stop(self):
        self.del_channel()

        # Break cycles and immediately disable the callback mechanism.
        self.read_handlers = ()
        self.write_handlers = ()

    def close(self):
        self.stop()
        self.socket.close()

    def readable(self):
        return self.read_handlers and self._readable()

    def writable(self):
        return self.write_handlers and self._writable()

    def handle_read(self):
        handlers = self.read_handlers
        if handlers:
            generator = handlers.pop(0)
            for more in generator:
                if more:
                    handlers.insert(0, generator)
                break

    def handle_write(self):
        handlers = self.write_handlers
        if handlers:
            generator = handlers.pop(0)
            for more in generator:
                if more:
                    handlers.insert(0, generator)
                break

    def handle_connect(self):
        pass

    def handle_accept(self):
        self.handle_read()

    def handle_close(self):
        pass                    # Don't close channel here.

    def handle_error(self):
        if not self.connected:
            # This is how asyncore handles "connection refused" in
            # Python 2.6.  See comment in __handle_connect() below.
            self.handle_write()
        else:
            logging.exception('Error in asyncore event handler:')

class Stream(_Stream):
    def _init(self, sock, aggressive_write=False, channel_factory=Channel):
        self.aggressive_write = aggressive_write

        self.channel = channel_factory(self.mgr.map, readable=self.readable)
        if isinstance(sock, tuple):
            self.channel.create_socket(*sock)
        elif isinstance(sock, socket.SocketType):
            sock.setblocking(0)
            self.channel.connected = True
            self.channel.set_socket(sock)
        else:
            raise TypeError

        self.channel.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)

    def set_rate_limit_read(self, rate, capacity):
        from messkit.token_bucket import TokenBucket
        self.token_bucket = TokenBucket(rate, capacity)

    def readable(self):
        if hasattr(self, 'token_bucket'):
            return self.token_bucket.authorize()
        else:
            return True

    def debit_read(self, bytes):
        if hasattr(self, 'token_bucket'):
            self.token_bucket.debit(bytes)

    def _close(self, flush, cb):
        if not flush or not self.connected:
            self.channel.close()
            self.issue_callback(cb)
        else:
            self.schedule_write_handler(self.__handle_close(cb))

    def __handle_close(self, cb):
        self.channel.close()
        invoke_callback(cb)
        yield False

    def _connect(self, address, cb):
        try:
            self.channel.connect(address)
        except socket.error, e: # E.g., error 64.
            self.issue_callback(cb, e)
        else:
            self.schedule_write_handler(self.__handle_connect(cb))

    def __handle_connect(self, cb):
        if not self.channel.connected:
            # The Python 2.6 version of asyncore swallows SO_ERROR and
            # raises socket.error, whereupon handle_error() is called.
            error = sys.exc_info()[1] or True
        else:
            # The Python 2.5 version of asyncore calls handle_connect()
            # and handle_write() without checking if the connection
            # attempt succeeded.
            errno = self.channel.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
            if errno == 0:
                error = None
            else:
                error = socket.error(errno, os.strerror(errno))
                self.error = error

        invoke_callback(cb, error)
        yield False

    def _flush(self, buf, cb=None):
        if not self.channel.write_handlers or cb is not None:
            self.schedule_write_handler(self.__handle_flush(buf, cb))

    def __handle_flush(self, buf, cb):
        while buf.size and not self.error:
            s = buf.peek(self.send_size)
            try:
                n = self.channel.send(s)
            except socket.error, e:
                self._signal_error(e)
                return

            if n:
                self.bytes_out += n
                buf.skip(n)

            if n < len(s) or (not self.aggressive_write and buf.size):
                yield True

        invoke_callback(cb)

    def _fill(self, buf, size, cb):
        start = self.bytes_in
        stop = start + size
        self.schedule_read_handler(self.__handle_fill(buf, start, stop, cb))

    def __handle_fill(self, buf, start, stop, cb):
        while self.bytes_in < stop and not self.eof_seen:
            try:
                s = self.channel.recv(self.recv_size)
            except socket.error, e:
                self._signal_error(e)
                return

            if not s:
                self.eof_seen = True
            else:
                buf.append(s)
                self.bytes_in += len(s)
                self.debit_read(len(s))

            if self.bytes_in < stop and not self.eof_seen:
                yield True

        invoke_callback(cb, self.bytes_in - start)

    def _get_socket(self):
        return self.channel.socket

    def schedule_read_handler(self, generator):
        self.channel.read_handlers.append(generator)

    def schedule_write_handler(self, generator):
        self.channel.write_handlers.append(generator)

class StreamServer(_StreamServer):
    def _init(self, params):
        self.channel = Channel(self.mgr.map)
        self.channel.create_socket(*params)
        self.channel.set_reuse_addr()

    def _close(self, cb):
        self.channel.close()
        self.issue_callback(cb)

    def _listen(self, address):
        self.channel.bind(address)
        self.channel.listen(5)

    def _accept(self, cb):
        self.schedule_read_handler(self.__handle_accept(cb))

    def __handle_accept(self, cb):
        result = self.channel.accept()

        while result is None:   # Probably never happens.
            yield True
            result = self.channel.accept()

        invoke_callback(cb, *result)
        yield False

    def schedule_read_handler(self, generator):
        self.channel.read_handlers.append(generator)

class StreamManager(_StreamManager):
    def __init__(self):
        _StreamManager.__init__(self, callback_callback=None)
        self.map = {}
        self.queue = Queue()

    def loop(self, **kwargs):
        while self.map:
            self.once(**kwargs)

    def once(self, **kwargs):
        while True:
            try:
                f = self.queue.get_nowait()
            except Empty:
                break

            invoke_callback(f)

        asyncore.loop(map=self.map, count=1, **kwargs)

    def _create_stream(self, *args):
        return Stream(self, *args)

    def _create_server(self, params):
        return StreamServer(self, params)

    def issue_main_thread_callback(self, f, *args):
        if f is not None:
            self.queue.put(lambda: f(*args))

def invoke_callback(f, *args):
    if f is not None:
        try:
            f(*args)
        except Exception:
            logging.exception('Error in callback:')
